
import pylab
from numpy import zeros, mgrid

import JT

#===============================================================================
#    Setup Registration
#===============================================================================
# Load in data from an xml file
data = Data.DataStructure()
data_path = r'C:\Users\bryan\bryan-code\trunk\Tests\new_data_structure.xml' 
data.load(data_path)

drr = JT.DRR()
reg=JT.MultiIterationRegistration()
# Set additional arguments for cost function
reg.drr.renWin.SetOffScreenRendering(1)
reg._plot = False
reg._debug = True
reg._logGradientResult = False
reg.SetIterations(4)

#===============================================================================
#    Setup Optimizer
#===============================================================================      
gd = RegularStepGradientDescentOptimizer()
reg.SetOptimizer(gd)
# Set the cost function
gd.SetCostFunction(reg)

reg.SetInitialPose(pose)
reg._initialStartingPose = pose

Volume = data.Tree.Frame[vert].Volume[0]
FixedImage = data.Tree.Frame[vert].FixedImage[proj]       

#===============================================================================
#    Setup DRR
#===============================================================================
# Load all of the data elements from the structure into the reg class
drr.LoadCalFile(data.get_path(FixedImage.Cal))
drr.LoadStaFile(data.get_path(Volume.STA))
drr.SetVolumeFileName(data.get_path(Volume))
drr.SetCenter(data.get_center(Volume, body))
drr.renWin.SetSize(512,512)

#Set the volume mapper to RayCast (more accurate, but slower)
#Comment out to use TextureMapper
drr.volume.UseRayCastMapper()
drr.InteractionOff()
drr.renWin.SetOffScreenRendering(1)

reg.SetVolume(data.Frames[vert].Volumes[0].data)
reg.SetFixedImage(data.Frames[vert].FixedImages[proj].data.GetImage())
#reg.SetRegionOfInterest(data.Frames[vert].FixedImages[proj].ROIs[body].data)
reg.SetRegionOfInterest(None)
reg.SetCenter(data.Frames[vert].Volumes[0].BodyCenters[body])
reg.SetImageMetric(method)


# Setup grid to evaluate over
xdim = 2
ydim = 2
m = mgrid[-xdim:xdim+1,-ydim:ydim+1]

gridx = m[0,:,:].flatten()
gridy = m[1,:,:].flatten()
scale = 1

val = zeros((2*xdim+1,2*ydim+1))
iter = 0

if 0:
    for i in range(-xdim,xdim+1):
        for j in range(-ydim,ydim+1):
#            pylab.imshow(currentImage)
            pose = [i*scale,j*scale,0,0,0,0]
            val[i][j] = reg.GetValue(pose)
            print "Iteration: %i [%i,%i], Metric Value: %f" % (iter,i*scale,j*scale,val[i][j])
            iter += 1
        
if 1:
    for i in range(-gridx/2,gridx/2):
        movingImage=fixedImage.rotate(i*2)
        val[i][j] = IM.GD(fixedImage, movingImage)
        print "Iteration: %i, Metric Value: %f" % (iter,val[i][j])
        iter += 1

pylab.figure

pylab.imshow(val)
